# 导入必要的库
import matplotlib.pyplot as plt
import pandas as pd
from math import sqrt

from keras import Sequential
from keras.layers import LSTM, Dense
from keras.regularizers import l2
from numpy import concatenate
from pandas import read_table
from pandas import DataFrame
from pandas import concat
from sklearn.preprocessing import MinMaxScaler, LabelEncoder
from sklearn.metrics import mean_squared_error
from pyemd import emd
#from pyhht.visualization import plot_imfs
from keras.models import load_model
import numpy as np

# 固定随机种子以复现结果
seed = 42
np.random.seed(seed)


# 把单变量的数据转化为有监督的数据
# convert series to supervised learning
def series_to_supervised(data, n_in=1, n_out=1, dropnan=True):
    n_vars = 1 if type(data) is list else data.shape[1]
    df = DataFrame(data)
    cols, names = list(), list()
    # input sequence (t-n, ... t-1)
    for i in range(n_in, 0, -1):
        cols.append(df.shift(i))
        names += [('var%d(t-%d)' % (j + 1, i)) for j in range(n_vars)]
    # forecast sequence (t, t+1, ... t+n)
    for i in range(0, n_out):
        cols.append(df.shift(-i))
        if i == 0:
            names += [('var%d(t)' % (j + 1)) for j in range(n_vars)]
        else:
            names += [('var%d(t+%d)' % (j + 1, i)) for j in range(n_vars)]
    # put it all together
    agg = concat(cols, axis=1)
    agg.columns = names
    # drop rows with NaN values
    if dropnan:
        agg.dropna(inplace=True)
    return agg


# 加载数据集
# load dataset
dataset = pd.read_csv('pollution.csv', header=0, index_col=0)
values = dataset.values
# integer encode direction
encoder = LabelEncoder()
values[:, 4] = encoder.fit_transform(values[:, 4])
# ensure all data is float
values = values.astype('float32')
# normalize features
scaler = MinMaxScaler(feature_range=(0, 1))
# scaler = StandardScaler()
scaled = scaler.fit_transform(values)
#emd = emd()

'''
进行emd分解
'''
pollution = values[:, 0]
imfs = emd(pollution)
#plot_imfs(pollution, np.array(imfs))
imfsValues = []
for imf in imfs:
    values[:, 0] = imf
    imfsValues.append(values.copy())
inv_yHats = []
inv_ys = []
for imf in imfsValues:
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled = scaler.fit_transform(imf)
    # scaleds.append(scaled)
    n_hours = 4
    n_features = 8
    reframed = series_to_supervised(scaled, n_hours, 1)
    values = reframed.values
    n_train_hours = 365 * 24 * 4
    train = values[:n_train_hours, :]
    test = values[n_train_hours:, :]

    # 监督学习结果划分,test_x.shape = (, 8)
    n_obs = n_hours * n_features
    train_x, train_y = train[:, :n_obs], train[:, -n_features]
    test_x, test_y = test[:, :n_obs], test[:, -n_features]

    # 为了在LSTM中应用该数据，需要将其格式转化为3D format，即[Samples, timesteps, features]
    train_X = train_x.reshape((train_x.shape[0], n_hours, n_features))
    test_X = test_x.reshape((test_x.shape[0], n_hours, n_features))

    model = Sequential()
    model.add(
        LSTM(20, input_shape=(train_X.shape[1], train_X.shape[2]), return_sequences=True, kernel_regularizer=l2(0.005),
             recurrent_regularizer=l2(0.005)))
    model.add(LSTM(20, kernel_regularizer=l2(0.005), recurrent_regularizer=l2(0.005)))
    model.add(Dense(1))
    model.compile(loss='mae', optimizer='adam')
    history = model.fit(train_X, train_y, epochs=500, batch_size=2 ** 10, validation_data=(test_X, test_y))

    # make the prediction,为了在原始数据的维度上计算损失，需要将数据转化为原来的范围再计算损失
    yHat = model.predict(test_X)
    y = model.predict(train_X)
    test_X = test_X.reshape((test_X.shape[0], n_hours * n_features))

    '''
        这里注意的是保持拼接后的数组  列数  需要与之前的保持一致
    '''
    inv_yHat = concatenate((yHat, test_X[:, -7:]), axis=1)  # 数组拼接
    inv_yHat = scaler.inverse_transform(inv_yHat)
    inv_yHat = inv_yHat[:, 0]
    inv_yHats.append(inv_yHat)

    test_y = test_y.reshape((len(test_y), 1))
    inv_y = concatenate((test_y, test_X[:, -7:]), axis=1)
    inv_y = scaler.inverse_transform(inv_y)  # 将标准化的数据转化为原来的范围
    inv_y = inv_y[:, 0]
    inv_ys.append(inv_y)

inv_yHats = np.array(inv_yHats)
inv_yHats = np.sum(inv_yHats, axis=0)
inv_ys = np.array(inv_ys)
inv_ys = np.sum(inv_ys, axis=0)
rmse = sqrt(mean_squared_error(inv_yHats, inv_ys))
print('Test RMSE: %.3f' % rmse)
inv_y = inv_ys[-24 * 3:]
inv_yHat = inv_yHats[-24 * 3:]
plt.plot(inv_yHat, label='forecast')
plt.plot(inv_y, label='observation')
plt.ylabel('pm2.5')
plt.legend()
plt.show()
# calculate RMSE
rmse = sqrt(mean_squared_error(inv_y, inv_yHat))
print('Test RMSE: %.3f' % rmse)

